'''
Created on Feb 1, 2012

@author: xzhu
'''
from CONF import CONF
from RES import RES
from IONISTATE import IONISTATE
from IONISTATE import combine_list

class PROT():
    '''
    MCCE protein type
    '''


    def __init__(self):
        '''
        Constructor
        '''
        self.residues   = []
        self.conformers = []
        self.iStates    = []       # possible ionization states of the residues in this protein
        self.selfEner   = 0.0
        self.pairEner   = 0.0
        self.total_ener = 0.0
                
    def init_from_res(self, residues):
        '''
        initialize the protein from some residues
        '''
        
        self.residues = residues
        for res in residues:
            for conf in res.conformers:
                self.conformers.append(conf)
        
    def load_from_files(self, head3 = 'head3.lst', fort38 = 'fort.38', energies = 'energies'):
        '''
        load all the information about the conformers of the protein from 'head3.lst', 'fort.38', 'energies'
        then divide all conformers into individual residues
        '''
        
        # load conformers from head3.lst
        lines = open(head3).readlines()
        lines.pop(0)
        #print lines[0]
        for line in lines:
            # print line
            newConf = CONF()
            newConf.head3_init(line)
            self.conformers.append(newConf)
            
        # load info from fort.38
        lines = open(fort38).readlines()
        lines.pop(0)
        for line in lines:
            for conf in self.conformers:
                conf.fort_init(line)

        # for non-dummy comformers, load pairwise interaction from opp file
        for conf in self.conformers:
            if conf.confID[3:5] != 'DM':
                conf.ener_opp_init(energies + '/' + conf.confID + '.opp')
            
            
            
        # get all the residues in protein from info of all the loaded conformers
        for conf in self.conformers:
            confFound = False
            for res in self.residues:
                if conf.is_in_res(res):
                    res.conformers.append(conf)
                    confFound = True
                    break
                
            if not confFound:
                newRes = RES()
                newRes.init_from_conf(conf)
                self.residues.append(newRes)
            
        # get the total self energy of a residue
        for res in self.residues:
            tot_En = 0.0
            #print 'res name ', res.resID
            for conf in res.conformers:
        #        print 'conf ', conf.confID, 'self ', conf.selfEner, 'focc ', conf.focc, 'energy ', conf.selfEner * conf.focc
                tot_En += conf.selfEner * conf.focc
            #print 'total residue energy ', tot_En 
                       
    def get_self_ener(self):
        ''' 
        get the self energy of this protein
        this is intended for fixed protein (occ from fort.38) ONLY
        '''
        
        # treat the self energy and pairwise energy seperately
        self.selfEner   = 0.0
        self.pairEner   = 0.0
        self.total_ener = 0.0
        
        for conf in self.conformers:
            if conf.confID[3:5] == 'DM': continue
            self.selfEner += conf.selfEner * conf.focc
            
            for other_conf in self.conformers:
                if other_conf.confID[3:5] == 'DM': continue
                #if abs(conf.tPair[other_conf.iConf] * other_conf.focc * conf.focc) > 0.5: 
                    #print 'interaction %s, %s, %.2f' % (conf.confID, other_conf.confID, conf.tPair[other_conf.iConf] * other_conf.focc * conf.focc)

                self.pairEner += conf.tPair[other_conf.iConf] * other_conf.focc * conf.focc
                
        self.total_ener = self.selfEner + self.pairEner * 0.5
        #print 'self ', self.selfEner, 'pair ', self.pairEner,  'total ', self.total_ener


    def get_poss_iStates(self):
        '''
        all the possible ionization states
        '''
        
        sg = []
        for res in self.residues:
            sg.append(res.poss_ioni_state())

        all_sgStates = combine_list(sg)
        #print all_sgStates
        
        for sgState in all_sgStates:
            self.iStates.append(IONISTATE(sgState, self.residues))
            # print self.iStates[-1].signs[0]
        
    def display_iStates(self, fixProt):
        '''
        doc
        '''
        
        for state in self.iStates:
            #print state.signs

            state.state_ener(fixProt)
            print state.display_state()
